package com.zelix.analy;

import com.arthas.client.ArthasClient;
import com.arthas.client.ArthasTool;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.zip.ZipEntry;

public class analyIdea {

    //缓存key,无法动态拦截key的可以设置缓存。动态也可以使用缓存。
    static HashMap<String, String> keyMap = new HashMap() {{
        put("com/b/a/b/TQ.<init>(Ljava/io/Reader;JLcom/b/a/b/jV;)V", "12240739205347");
        put("com/b/a/a/CX.a(J)V", "9954405685380");
        put("com/b/a/a/CX.a(Ljava/util/List;J)Lkotlin/Pair;", "117118170400667");
    }};

    public static void main(String[] args) throws IOException {

        JarFile jarFile = null;
        try {
            jarFile = new JarFile("D:\\jetbrains\\idea\\lib\\app.jar");
        } catch (IOException e) {
            e.printStackTrace();
        }

        //远程调试
        //ArthasClient client = ArthasTool.getClient("192.168.228.128", 3658);
        //本地调试
        ArthasClient client = ArthasTool.getClient();

        if (jarFile == null || client == null) {
            return;
        }
        String rule = "";
        //Code With Me server
        // jarFile = new JarFile("D:\\tool\\lobby-server-fatjar.jar");
        rule = "com/b/a/a/CX";
        rule = "com/b/a/b/TQ";

        //IntelliJ IDEA
        //jarFile = new JarFile("D:\\jetbrains\\idea\\lib\\idea.jar");
        rule = "com/jetbrains/y/s/o";//这个忘记哪个版本了
        rule = "com/jetbrains/J/Q/D";//IntelliJ IDEA 2021.3.2 (Ultimate Edition)
        //IntelliJ IDEA 2022.1
        //jarFile = new JarFile("D:\\jetbrains\\idea\\lib\\app.jar");
        rule = "com/jetbrains/x/z/z";//IntelliJ IDEA 2022.1 (Ultimate Edition) 证书链验证
        //rule="com/intellij/ide/x/c/F";//注册code分析
        //rule="com/intellij/ide/x/c/zJ";//json解析
        //IntelliJ IDEA 2022.2
        //jarFile = new JarFile("D:\\jetbrains\\idea\\lib\\app.jar");
        rule = "com/jetbrains/C/x/C";//IntelliJ IDEA 2022.2 (Ultimate Edition) 证书链验证
        //rule="com/jetbrains/C/x/l";//证书加载以及getEncoded验证
        //rule="com/intellij/ide/x/x/e";//json解析
        rule = "com/intellij/ide/x/C/P";//注册面板


        //方法太多时可以过滤关注的方法,名称和签名更精确
        ArrayList<String> method = new ArrayList<>();
        //打开注册页面触发,通过button.discover.server找到按钮对应的类的属性
        method.add("<init>(Lcom/intellij/ide/x/a;SLcom/intellij/ide/x/C/xP;ILcom/intellij/ide/x/N/I;SLjava/lang/String;)V");
        //点击按钮触发,java.util.EventObject.getSource()比对找到对应的button.discover.server执行方法
        method.add("actionPerformed(Ljava/awt/event/ActionEvent;)V");
        //点击“Discover Server”按钮触发
        method.add("x(SISZ)V");
        //点击“Discover Server”按钮触发
        method.add("x(Ljava/util/concurrent/atomic/AtomicBoolean;Lcom/intellij/openapi/util/Computable;JLjava/util/function/Consumer;JZ)V");
        method.add("x(Lcom/intellij/openapi/util/Computable;)Ljava/lang/Object;");

        Enumeration<JarEntry> entries = jarFile.entries();
        while (entries.hasMoreElements()) {
            JarEntry entry = entries.nextElement();
            if (entry.getName().endsWith(".class") && entry.getName().startsWith(rule)) {//只分析class文件
                //去除文件名后缀 ".class"
                analy(client, jarFile, entry.getName().substring(0, entry.getName().length() - 6), method);
            }
        }

        client.disconnect();
    }

    public static void analy(ArthasClient client, JarFile jarFile, String cls, ArrayList<String> method) throws IOException {

        //key为null时拦截还是跳过。无法拦截时设置为true,关键值存入keyMap。可以拦截时设置为false
        boolean skip = false;

        //获取类所在的classloader,idea中的一些方法不加classloader去运行,可能找不到类
        String classLoader = ArthasTool.getClassLoader(client, cls);
        if (classLoader == null) {
            System.out.println("类 " + cls + " 还未加载,请点击相应功能触发其加载.");
            return;
        }

        /**
         * 两种获取计算key值的方式,一种是通过字符串加密函数,一种是通过获取反射方法的方法,以防解析的方法中只有一种方式
         */
        //字符串解密的ognl表示方式
        String stringMethodFormat = null;
        //监听字符串解密方法
        String watchStringMethodFormat = null;
        //获取反射Method的ognl表示方式
        String reflectMethodFormat = "ognl '@%s@%s(%sL,%sL).toString' -c " + classLoader;
        //获取asm描述Method的ognl表示方式
        String asmMethodFormat = "ognl '#method=@%s@%s(%sL,%sL),#des=@jdk.internal.org.objectweb.asm.Type@getType(#method),#ower=@jdk.internal.org.objectweb.asm.Type@getType(#method.getDeclaringClass()),#ower+\":\"+#method.getName()+\":\"+#des' -c " + classLoader;
        //获取asm描述Field的ognl表示方式
        String asmFieldFormat = "ognl '#field=@%s@%s(%sL,%sL),#des=@jdk.internal.org.objectweb.asm.Type@getType(#field.getType()),#ower=@jdk.internal.org.objectweb.asm.Type@getType(#field.getDeclaringClass()),#ower+\":\"+#field.getName()+\":\"+#des' -c " + classLoader;
        //监听获取反射Method的方法
        String watchReflectMethodFormat = "watch %s %s {params[1]} 'params.length==2 && params[0] instanceof java.lang.Long && params[1] instanceof java.lang.Long && returnObj !=null && returnObj instanceof java.lang.reflect.Method && params[0]==%sL' -n 1";
        //监听获取反射Field的方法
        String watchReflectFieldFormat = "watch %s %s {params[1]} 'params.length==2 && params[0] instanceof java.lang.Long && params[1] instanceof java.lang.Long && returnObj !=null && returnObj instanceof java.lang.reflect.Field && params[0]==%sL' -n 1";

        ZipEntry entry = jarFile.getEntry(cls + ".class");

        if (entry != null) {
            ClassReader cr = new ClassReader(jarFile.getInputStream(entry));
            ClassNode cn = new ClassNode();
            cr.accept(cn, 0);

            List<MethodNode> methods = cn.methods;

            //查找符合本类的字符串解密方法并初始化解密方法,字符串加密都在本类
            for (MethodNode methodNode : methods) {
                if (methodNode.desc.equals("(IJ)Ljava/lang/String;")) {
                    stringMethodFormat = "ognl '@" + cls.replaceAll("/", ".") + "@" + methodNode.name + "(%d,%sL^%sL)' -c " + classLoader;
                    watchStringMethodFormat = "watch " + cls.replaceAll("/", ".") + " " + methodNode.name + " {params[1]} 'params.length==2 && params[0] instanceof java.lang.Integer && params[1] instanceof java.lang.Long && returnObj !=null && returnObj instanceof String && params[0]==%d' -n 1";
                }
            }

            System.out.println("************" + cls + "************");

            //遍历方法
            for (MethodNode methodNode : methods) {

                //关注的方法有数据,跳过非关注的方法
                if (method.size() > 0 && !method.contains(methodNode.name + methodNode.desc)) {
                    continue;
                }

                System.out.println("------------" + methodNode.name + methodNode.desc + "------------");
                //计算long值,遍历每个方法前初始化,不同方法值不同
                String key = keyMap.get(cls + "." + methodNode.name + methodNode.desc);
                InsnList insnList = methodNode.instructions;
                //遍历指令
                for (int i = 0; i < insnList.size(); i++) {
                    /**
                     * SIPUSH 3671
                     * LDC 985935450605195468L
                     * LLOAD 4
                     * LXOR
                     * INVOKEDYNAMIC n (IJ)Ljava/lang/String; handle[H_INVOKESTATIC com/jetbrains/y/s/o.a(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;] args[]
                     */
                    //查找invokedynamic指令
                    if (insnList.get(i) instanceof InvokeDynamicInsnNode) {
                        InvokeDynamicInsnNode dynamicNode = (InvokeDynamicInsnNode) insnList.get(i);
                        int depth = 1;
                        //是不是解密字符方法签名,是不是调用本类内的方法,是的话一般是字符串加密用途,最后判断一下本类找没找到解密方法
                        if (dynamicNode.desc.equals("(IJ)Ljava/lang/String;") && dynamicNode.bsm.getOwner().equals(cls) && stringMethodFormat != null) {
                            //不用判断i是否够减,InvokeDynamicInsnNode指令一般需要很多参数,放心使用
                            //根据情况加了一个深度判断
                            while (!(insnList.get(i - depth) instanceof InsnNode && insnList.get(i - depth).getOpcode() == Opcodes.LXOR) && depth < 10) {
                                depth++;
                            }
                            //可能depth深度跳出的循环,还要有判断
                            if (insnList.get(i - depth) instanceof InsnNode && insnList.get(i - depth).getOpcode() == Opcodes.LXOR) {//判断上一个是不是LOR
                                if (insnList.get(i - depth - 1) instanceof VarInsnNode && insnList.get(i - depth - 1).getOpcode() == Opcodes.LLOAD) {
                                    if (insnList.get(i - depth - 2) instanceof LdcInsnNode) {
                                        if (insnList.get(i - depth - 3) instanceof IntInsnNode) {
                                            LdcInsnNode ldcNode = (LdcInsnNode) insnList.get(i - depth - 2);
                                            IntInsnNode intNode = (IntInsnNode) insnList.get(i - depth - 3);

                                            //key未初始化,通过拦截字符串解密方法得到值
                                            if (!skip && key == null) {
                                                System.out.println("(String)请点击响应功能触发方法调用.");
                                                client.setTimeOut(2 * 60 * 1000L);
                                                String cmd = String.format(watchStringMethodFormat, intNode.operand);
                                                System.out.println(cmd);
                                                String rtn = client.sendCommand(cmd);
                                                key = getWatchReturnLong(rtn);
                                                if (key != null) {
                                                    //拦截到的是异或运算的,求解原值
                                                    key = String.valueOf(Long.valueOf(ldcNode.cst.toString()) ^ Long.valueOf(key));
                                                }
                                            }

                                            if (key != null) {
                                                String rtn = client.sendCommand(String.format(stringMethodFormat, intNode.operand, ldcNode.cst.toString(), key));
                                                System.out.println("字符串解密:(" + intNode.operand + "," + ldcNode.cst + ")" + rtn);
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                    /**
                     * LDC -4952167813042798704L
                     * LLOAD 4
                     * INVOKESTATIC a/b/gs.d(JJ)Ljava/lang/reflect/Method;
                     */
                    //MethodInsnNode
                    if (insnList.get(i) instanceof MethodInsnNode) {
                        MethodInsnNode methodInsnNode = (MethodInsnNode) insnList.get(i);
                        if (methodInsnNode.desc.equals("(JJ)Ljava/lang/reflect/Method;") || methodInsnNode.desc.equals("(JJ)Ljava/lang/reflect/Field;")) {
                            String methoType = "";
                            if (methodInsnNode.desc.equals("(JJ)Ljava/lang/reflect/Method;")) {
                                methoType = "Method";
                            } else {
                                methoType = "Field";
                            }
                            if (insnList.get(i - 1) instanceof VarInsnNode && insnList.get(i - 1).getOpcode() == Opcodes.LLOAD) {
                                if (insnList.get(i - 2) instanceof LdcInsnNode) {
                                    LdcInsnNode ldcNode = (LdcInsnNode) insnList.get(i - 2);
                                    if (!skip && key == null) {
                                        System.out.println("(" + methoType + ")请点击响应功能触发方法调用.");
                                        client.setTimeOut(2 * 60 * 1000L);
                                        String format = methoType.equals("Method") ? watchReflectMethodFormat : watchReflectFieldFormat;
                                        String cmd = String.format(format, methodInsnNode.owner.replaceAll("/", "."), methodInsnNode.name, ldcNode.cst.toString());
                                        System.out.println(cmd);
                                        String rtn = client.sendCommand(cmd);
                                        key = getWatchReturnLong(rtn);
                                    }
                                    if (key != null) {
                                        String rtn = client.sendCommand(String.format(reflectMethodFormat, methodInsnNode.owner.replaceAll("/", "."), methodInsnNode.name, ldcNode.cst.toString(), key));
                                        System.out.println((methoType.equals("Method") ? "方法" : "属性") + "解密:(" + ldcNode.cst + "," + key + ")" + rtn);
                                        //String format = methoType.equals("Method") ? asmMethodFormat : asmFieldFormat;
                                        //rtn = client.sendCommand(String.format(format, methodInsnNode.owner.replaceAll("/", "."), methodInsnNode.name, ldcNode.cst.toString(), key));
                                        //System.out.println("ASM对象描述" + asmdes(rtn));
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    public static String getWatchReturnLong(String rtn) {
        //超时时会返回null
        if (rtn == null) {
            return null;
        }
        String tag = "@Long[";
        int start = rtn.indexOf(tag);
        int end = rtn.indexOf("]", start);
        return rtn.substring(start + tag.length(), end);
    }

    /**
     * 返回ower+name+des的形式
     *
     * @param rtn
     * @return
     */
    public static String asmdes(String rtn) {
        rtn = ArthasTool.removeStringLebal(rtn);
        String[] split = rtn.split(":");
        if (split.length == 3) {
            if (split[0].startsWith("L") && split[0].endsWith(";")) {//owner应该是对象格式
                return split[0].substring(1, split[0].length() - 1) + ":" + split[1] + ":" + split[2];
            }
        }
        return null;
    }
}
